package com.ai.zuochengyun.phase01.class05;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Stack;

/**
 * 并查集
 */
public class Code01_UnionFind {

    public static void main(String[] args) {

    }

    public static class UnionSet<V> {
        public Map<V, Node<V>> nodes;

        public Map<Node<V>, Node<V>> parents;

        /**
         * 代表节点才在此集合中存在
         */
        public Map<Node<V>, Integer> sizes;

        public UnionSet(List<V> values) {
            this.nodes = new HashMap<>();
            this.parents = new HashMap<>();
            this.sizes = new HashMap<>();

            for (V v : values) {
                Node<V> cur = new Node<>(v);
                nodes.put(v, cur);
                parents.put(cur, cur);
                sizes.put(cur, 1);
            }
        }

        /**
         * 找到原始的祖先节点
         * @param cur
         * @return
         */
        public Node<V> find(Node<V> cur) {
            Stack<Node<V>> stack = new Stack<>();
            while (cur != parents.get(cur)) {
                stack.push(cur);
                cur = parents.get(cur);
            }

            // 将沿途所有的节点的祖先节点修改为了当前祖先
            while (!stack.isEmpty()) {
                parents.put(stack.pop(), cur);
            }
            return cur;
        }

        public boolean isSameSet(V a, V b) {
            return find(nodes.get(a)) == find(nodes.get(b));
        }

        public void union(V a, V b) {
            Node<V> aHead = find(nodes.get(a));
            Node<V> bHead = find(nodes.get(b));
            // 如果aHead与bHead相等，代表在一个集合中，那么不用做任何处理
            if (aHead != bHead) {
                // 找到哪个头节点挂的节点较多
                Integer aSize = sizes.get(aHead);
                Integer bSize = sizes.get(bHead);
                Node<V> big = aSize >= bSize ? aHead : bHead;
                Node<V> small = big == aHead ? bHead : aHead;

                // 将节点较小的头挂到节点较多的头上
                parents.put(small, big);
                sizes.put(big, aSize + bSize);
                sizes.remove(small);
            }
        }
    }

    private static class Node<V> {
        V value;

        public Node(V value) {
            this.value = value;
        }
    }
}
